#include "hypercube_hmm_mink.hpp"
#include "helib_atm_wrapper.hpp"

static inline double log2_noise_bound(const helib::Ctxt &ctxt) {
  const double ln2 = NTL::log(NTL::xdouble(2l));
  return NTL::log(ctxt.totalNoiseBound()) / ln2;
}

class hmm_mink_ch : public constant_handle {
public:
  hmm_mink_ch(const helib::EncryptedArray &ea, const cube_signature &eh_mnk)
      : eh_m_(eh_mnk.size(0 /* m */)), eh_n_(eh_mnk.size(1 /* n */)),
        min_(eh_mnk.size(2 /* k */)), constant_handle(ea) {
    ASSERT_PRINTF(ea_.dimension() == 2,
                  "Only support two-dimensional hypercube\n");
    if (eh_mnk.size(0 /* m */) != eh_mnk.size(2 /* k */)) {
      clean_left_cube_.init({eh_mnk.size(0 /* m */), eh_mnk.size(2 /* k */)});
    }
    if (eh_mnk.size(1 /* n */) != eh_mnk.size(2 /* k */)) {
      clean_right_cube_.init({eh_mnk.size(2 /* k */), eh_mnk.size(1 /* n */)});
    }
  }

  void init(hmm_metric_collector *collector) override;

  bool clean_left() const { return !clean_left_cube_.empty(); }
  bool clean_right() const { return !clean_right_cube_.empty(); }

private:
  long eh_m_;
  long eh_n_;
  long min_;
  cube_signature clean_left_cube_;
  cube_signature clean_right_cube_;

public:
  std::array<helib::zzX, 2> rr_clean_;
  std::vector<helib::zzX> ra_lefts_;
  std::vector<helib::zzX> ra_rights_;
  std::array<helib::zzX, 2> ra_corrects_;
  std::vector<std::array<helib::zzX, 2>> sc_corrects_;
};

void hmm_mink_ch::init(hmm_metric_collector *collector) {
  auto hmm_mink = static_cast<hmm_mink_collector *>(collector);
  ASSERT_PRINTF(hmm_mink != nullptr,
                "Failed to static_cast to `hmm_mink_collector *`\n");
  const long h_nrow = ea_.sizeOfDimension(0); //   m-dim
  const long h_ncol = ea_.sizeOfDimension(1); //   n-dim
  // const long eh_min = std::min(eh_m_, eh_n_);
  const long eh_min = min_;
  const auto threads = std::min(threads_, static_cast<int>(h_nrow));

  hmm_mink->start_encoding_ra(); //! Encoding in RotateAlign - Start
  ra_lefts_.resize(eh_min);
  ra_rights_.resize(eh_min);
  std::vector<NTL::zz_pX> col_i(eh_n_);
  std::vector<NTL::zz_pX> lrow_i(eh_min);
  std::vector<NTL::zz_pX> rrow_i(eh_min);
#pragma omp parallel
#pragma omp single
  {
    // align A
#pragma omp taskloop nogroup
    for (long i = 0; i < eh_min; i++) {
      auto tid = omp_get_thread_num();
      hmm_mink->start_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Start
      std::vector<long> slots(ea_.size(), 0);
      for (long j = 0; j < h_ncol; j++) {
        for (long k = 0; k * eh_min + i < h_nrow; k++) {
          slots[(k * eh_min + i) * h_ncol + j] = 1l;
        }
      }
      hmm_mink->stop_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Completed

      hmm_mink->start_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Start
      ea_.encode(ra_lefts_[i], slots);
      hmm_mink->stop_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Completed
    }

    // align B
#pragma taskloop nogroup
    for (long i = 0; i < eh_n_; i++) {
      auto tid = omp_get_thread_num();
      NTL::zz_pPush push(ea_.getPAlgebra().getP());

      hmm_mink->start_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Start
      std::vector<long> slots(ea_.size(), 0);
      for (long j = 0; j < h_nrow; j++) {
        slots[j * h_ncol + i] = 1l;
      }
      hmm_mink->stop_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Completed

      hmm_mink->start_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Start
      helib::zzX poly;
      ea_.encode(poly, slots);
      helib::convert(col_i[i], poly);
      hmm_mink->stop_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Completed
    }

    //! Encoding in RotateAlign (Correct) - Start
#pragma omp taskloop nogroup
    for (long i = 0; i < eh_min; i++) {
      auto tid = omp_get_thread_num();
      //!! `zz_pPush` must have a copy in each thread
      NTL::zz_pPush push(ea_.getPAlgebra().getP());

      hmm_mink->start_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Start
      std::vector<long> slots1(ea_.size(), 0);
      std::vector<long> slots2(ea_.size(), 0);
      for (long j = 0; j < h_ncol; j++) {
        for (long k = 0; k * eh_min + i < h_nrow; k++) {
          if (i + j >= h_ncol) {
            slots1[(k * eh_min + i) * h_ncol + j] = 1l;
          } else if (i + j < eh_n_) {
            slots2[(k * eh_min + i) * h_ncol + j] = 1l;
          }
        }
      }
      hmm_mink->stop_encoding_ra_raw2vec(tid); //! Raw 2 Vec - Completed

      hmm_mink->start_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Start
      helib::zzX rrow, lrow;
      ea_.encode(rrow, slots1);
      ea_.encode(lrow, slots2);
      helib::convert(rrow_i[i], rrow);
      helib::convert(lrow_i[i], lrow);
      hmm_mink->stop_encoding_ra_vec2ptxt(tid); //! Vec 2 Ptxt - Completed
    }
  }
  hmm_mink->start_encoding_ra_re_ptxt(); //! Re/Ptxt - Start
  NTL::zz_pPush push(ea_.getPAlgebra().getP());
  NTL::zz_pX L, R;
  for (long i = 0; i < eh_min; i++) {
    L += lrow_i[i];
    R += rrow_i[i];
  }
  ra_corrects_[0] = helib::balanced_zzX(R);
  ra_corrects_[1] = helib::balanced_zzX(L);
  hmm_mink->stop_encoding_ra_re_ptxt(); //! Re/Ptxt - Stop
  //! Encoding in RotateAlign (Correct) - Completed

  hmm_mink->start_encoding_ra_re_ptxt(); //! Re/Ptxt - Start
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
  for (long j = 0; j < eh_min; j++) {
    NTL::zz_pPush push(ea_.getPAlgebra().getP());
    NTL::zz_pX res;
    for (long k = 0; k * eh_min + j < eh_n_; k++) {
      res += col_i[k * eh_min + j];
    }
    ra_rights_[j] = helib::balanced_zzX(res);
  }
  hmm_mink->stop_encoding_ra_re_ptxt(); //! Re/Ptxt - Stop
  hmm_mink->stop_encoding_ra();         //! Encoding in RotateAlign - Completed

  hmm_mink->start_encoding_sc();         //! Encoding in ShiftCompute - Start
  hmm_mink->start_encoding_sc_re_ptxt(); //! Re/Ptxt - Start
  //* we compute on "encoded" plaintexts to avoid encoding full data
  sc_corrects_.resize(min_);
  // Using "after encoded" plaintexts to construct targets
  //* h_nrow^2 * \phi(m) times NTL::AddMod
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
  for (long i = 1; i < min_; i++) {
    NTL::zz_pPush push(ea_.getPAlgebra().getP());
    NTL::zz_pX L, R;
    for (long j = 0; j < eh_n_; j++) {
      j < eh_n_ - i ? L += col_i[j] : R += col_i[j];
    }
    sc_corrects_[i][0] = helib::balanced_zzX(L);
    sc_corrects_[i][1] = helib::balanced_zzX(R);
  }
  hmm_mink->stop_encoding_sc_re_ptxt(); //! Re/Ptxt - Stop
  hmm_mink->stop_encoding_sc();         //! Encoding in ShiftCompute - Completed

  if (!clean_left_cube_.empty()) {
    std::vector<NTL::zz_pX> row_i(clean_left_cube_.size(0));
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
    for (auto i = 0; i < clean_left_cube_.size(0); i++) {
      NTL::zz_pPush push(ea_.getPAlgebra().getP());
      std::vector<long> slots(ea_.size(), 0);
      for (int j = 0; j < clean_left_cube_.size(1); j++) {
        slots[i * h_ncol + j] = 1l;
      }
      helib::zzX row;
      ea_.encode(row, slots);
      helib::convert(row_i[i], row);
    }
    NTL::zz_pX res;
    NTL::zz_pPush push(ea_.getPAlgebra().getP());
    for (auto i = 0; i < clean_left_cube_.size(0); i++) {
      res += row_i[i];
    }
    rr_clean_[0] = helib::balanced_zzX(res);
  }
  if (!clean_right_cube_.empty()) {
    std::vector<NTL::zz_pX> row_i(clean_right_cube_.size(0));
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
    for (auto i = 0; i < clean_right_cube_.size(0); i++) {
      NTL::zz_pPush push(ea_.getPAlgebra().getP());
      std::vector<long> slots(ea_.size(), 0);
      for (int j = 0; j < clean_right_cube_.size(1); j++) {
        slots[i * h_ncol + j] = 1l;
      }
      helib::zzX row;
      ea_.encode(row, slots);
      helib::convert(row_i[i], row);
    }
    NTL::zz_pX res;
    NTL::zz_pPush push(ea_.getPAlgebra().getP());
    for (auto i = 0; i < clean_right_cube_.size(0); i++) {
      res += row_i[i];
    }
    rr_clean_[1] = helib::balanced_zzX(res);
  }
}

class hhmm_mink_constant_multipler : public hmm_constant_multipler {
public:
  hhmm_mink_constant_multipler(const helib::EncryptedArray &ea,
                               hmm_metric_collector *collector,
                               const cube_signature &eh_mnk)
      : ea_(ea), eh_col_(eh_mnk.size(1 /* n */)), min_(eh_mnk.size(2 /* k */)),
        constant_handle_(new hmm_mink_ch(ea, eh_mnk)) {
    collector_ = static_cast<hmm_mink_collector *>(collector);
    ASSERT_PRINTF(collector_ != nullptr,
                  "Failed to static_cast to `hmm_mink_collector *`\n");

    constant_handle_->init(collector_);
  }

  ~hhmm_mink_constant_multipler() override { delete constant_handle_; }

  void rotate_align_mask_left(helib::Ctxt &ctxt, std::size_t i) {
    if (no_encoding_reuse_) {
      collector_->start_no_encoding_reuse(omp_get_thread_num());
      const long h_nrow = ea_.sizeOfDimension(0);
      const long h_ncol = ea_.sizeOfDimension(1);
      helib::zzX poly;
      std::vector<long> slots(ea_.size(), 0);
      for (long j = 0; j < h_ncol; j++) {
        for (long k = 0; k * min_ + i < h_nrow; k++) {
          slots[(k * min_ + i) * h_ncol + j] = 1l;
        }
      }
      ea_.encode(poly, slots);
      collector_->stop_no_encoding_reuse(omp_get_thread_num());
      ctxt.multByConstant(poly);
      return;
    }
    ctxt.multByConstant(constant_handle_->ra_lefts_[i]);
  }

  void rotate_align_mask_right(helib::Ctxt &ctxt, std::size_t i) {
    if (no_encoding_reuse_) {
      collector_->start_no_encoding_reuse(omp_get_thread_num());
      const long h_nrow = ea_.sizeOfDimension(0);
      const long h_ncol = ea_.sizeOfDimension(1);
      helib::zzX poly;
      std::vector<long> slots(ea_.size(), 0);
      for (long k = 0; k * min_ + i < eh_col_; k++) {
        for (long j = 0; j < h_nrow; j++) {
          slots[j * h_ncol + (k * min_ + i)] = 1l;
        }
      }
      ea_.encode(poly, slots);
      collector_->stop_no_encoding_reuse(omp_get_thread_num());
      ctxt.multByConstant(poly);
      return;
    }
    ctxt.multByConstant(constant_handle_->ra_rights_[i]);
  }

  void rotate_align_correct_mask(helib::Ctxt &c1, helib::Ctxt &c2) {
    if (no_encoding_reuse_) {
      collector_->start_no_encoding_reuse(omp_get_thread_num());
      const long h_nrow = ea_.sizeOfDimension(0);
      const long h_ncol = ea_.sizeOfDimension(1);
      std::vector<long> slots1(ea_.size(), 0);
      std::vector<long> slots2(ea_.size(), 0);
      for (long i = 0; i < min_; i++) {
        for (long j = 0; j < h_ncol; j++) {
          for (long k = 0; k * min_ + i < h_nrow; k++) {
            if (i + j >= h_ncol) {
              slots1[(k * min_ + i) * h_ncol + j] = 1l;
            } else if (i + j < eh_col_) {
              slots2[(k * min_ + i) * h_ncol + j] = 1l;
            }
          }
        }
      }
      helib::zzX poly1, poly2;
      ea_.encode(poly1, slots1);
      ea_.encode(poly2, slots2);
      collector_->stop_no_encoding_reuse(omp_get_thread_num());
      c1.multByConstant(poly1);
      c2.multByConstant(poly2);
      return;
    }
    c1.multByConstant(constant_handle_->ra_corrects_[0]);
    c2.multByConstant(constant_handle_->ra_corrects_[1]);
  }

  void shift_compute_correct_mask(helib::Ctxt &c1, helib::Ctxt &c2,
                                  std::size_t i) {
    if (no_encoding_reuse_) {
      collector_->start_no_encoding_reuse(omp_get_thread_num());
      const long h_nrow = ea_.sizeOfDimension(0);
      const long h_ncol = ea_.sizeOfDimension(1);
      std::vector<long> slots1(ea_.size(), 0);
      std::vector<long> slots2(ea_.size(), 0);
      for (long j = 0; j < eh_col_; j++) {
        for (long k = 0; k < h_nrow; k++) {
          j < eh_col_ - i ? slots1[k * h_ncol + j] = 1l
                          : slots2[k * h_ncol + j] = 1l;
        }
      }
      helib::zzX poly1, poly2;
      ea_.encode(poly1, slots1);
      ea_.encode(poly2, slots2);
      collector_->stop_no_encoding_reuse(omp_get_thread_num());
      c1.multByConstant(poly1);
      c2.multByConstant(poly2);
      return;
    }
    c1.multByConstant(constant_handle_->sc_corrects_[i][0]);
    c2.multByConstant(constant_handle_->sc_corrects_[i][1]);
  }

  void clean_mask_left(helib::Ctxt &ctxt) {
    if (clean_left_ && constant_handle_->clean_left()) {
      ctxt.multByConstant(constant_handle_->rr_clean_[0]);
    }
  }

  void clean_mask_right(helib::Ctxt &ctxt) {
    if (clean_right_ && constant_handle_->clean_right()) {
      ctxt.multByConstant(constant_handle_->rr_clean_[1]);
    }
  }

  void set_no_encoding_reuse(bool val) override { no_encoding_reuse_ = val; }
  void set_clean_left() { clean_left_ = true; }
  void set_clean_right() { clean_right_ = true; }

  long min_dim() const { return min_; }
  long eh_col() const { return eh_col_; }

private:
  long eh_col_;
  long min_;

  bool clean_left_ = false;
  bool clean_right_ = false;

  const helib::EncryptedArray &ea_;
  hmm_mink_collector *collector_ = nullptr;
  hmm_mink_ch *constant_handle_ = nullptr;
};

hmm_constant_multipler *
build_hhmm_mink_constant_multipler(hmm_metric_collector *collector,
                                   const helib::EncryptedArray &ea,
                                   const cube_signature &eh_mnk) {
  return new hhmm_mink_constant_multipler(ea, collector, eh_mnk);
}

void set_hhmm_mink_constant_multipler_clean_left(hmm_constant_multipler *cm) {
  static_cast<hhmm_mink_constant_multipler *>(cm)->set_clean_left();
}

void set_hhmm_mink_constant_multipler_clean_right(hmm_constant_multipler *cm) {
  static_cast<hhmm_mink_constant_multipler *>(cm)->set_clean_right();
}

/********************************************************************/
/*** Homomorphic Square Matrix Multiplication Based On Hypercube  ***/
/********************************************************************/

static inline void replicate1d(hmm_mink_collector *collector,
                               const helib::EncryptedArray &ea,
                               helib::Ctxt &ctxt, long blockSize, long d) {
  long dSize = ea.sizeOfDimension(d);

  long sz = dSize / blockSize; // how many blocks fit in this dimension

  if (sz == 1)
    return; // nothing to do, only one block in this dimension

  long k = NTL::NumBits(sz);
  long e = 1;
  helib::Ctxt ctxt_orig = ctxt;

  // do the actual replication using "shift and add"
  for (long j = k - 2; j >= 0; j--) {
    helib::Ctxt tmp = ctxt;
    if (collector) {
      collector->start_computing_rr_rot();
    }
    ea.rotate1D(tmp, d, e * blockSize, true);
    if (collector) {
      collector->stop_computing_rr_rot();
    }

    if (collector) {
      collector->start_computing_rr_add();
    }
    ctxt += tmp;
    if (collector) {
      collector->stop_computing_rr_add();
    }

    // e -> 2*e
    e = 2 * e;

    long b = NTL::bit(sz, j); // bit j of sz
    // e -> e+b
    if (b) {
      if (collector) {
        collector->start_computing_rr_rot();
      }
      ea.rotate1D(ctxt, d, 1l * blockSize, true);
      if (collector) {
        collector->stop_computing_rr_rot();
      }

      if (collector) {
        collector->start_computing_rr_add();
      }
      ctxt += ctxt_orig;
      if (collector) {
        collector->stop_computing_rr_add();
      }
      e++;
    }
  }
}

// Helper inline function for correcting RotateAlign
inline void correct_rotate_align_A(hmm_mink_collector *hmm_mink,
                                   hhmm_mink_constant_multipler *hmm_mink_cm,
                                   const helib::EncryptedArray &ea,
                                   const long h_nrow, const long h_ncol,
                                   helib::Ctxt &A) {
  if (hmm_mink) {
    hmm_mink->start_computing_ra_scale(); //! RA.Scale - Start
  }
  helib::Ctxt tmp = A;
  hmm_mink_cm->rotate_align_correct_mask(tmp, A);
  if (hmm_mink) {
    hmm_mink->stop_computing_ra_scale(); //! RA.Scale - Stop
  }

  if (hmm_mink) {
    hmm_mink->start_computing_ra_rot(); //! RA.Rot - Start
  }
  ea.rotate1D(tmp, 1, -(h_ncol - h_nrow));
  if (hmm_mink) {
    hmm_mink->stop_computing_ra_rot(); //! RA.Rot - Stop
  }

  if (hmm_mink) {
    hmm_mink->start_computing_ra_add(); //! RA.Add - Start
  }
  A += tmp;
  if (hmm_mink) {
    hmm_mink->stop_computing_ra_add(); //! RA.Add - Stop
  }
}

// Helper inline function for ShiftCompute
inline void shift_compute(hmm_mink_collector *hmm_mink,
                          hhmm_mink_constant_multipler *hmm_mink_cm,
                          const helib::EncryptedArray &ea, const long h_nrow,
                          const helib::Ctxt &Ap, const helib::Ctxt &Bp,
                          helib::Ctxt &AB) {
  hmm_mink->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
  hmm_mink->start_computing_sc();
  hmm_mink->start_computing_sc_rot(); //! SC.Rot - Start
  helib::GeneralAutomorphPrecon_FULL hoisted_Ap(Ap, 1, ea);
  helib::GeneralAutomorphPrecon_FULL hoisted_Bp(Bp, 0, ea);
  hmm_mink->stop_computing_sc_rot(); //! SC.Rot - Stop
  for (int i = 0; i < hmm_mink_cm->min_dim(); i++) {
    hmm_mink->start_computing_sc_rot(); //! SC.Rot - Start
    auto A_i = hoisted_Ap.automorph(-i);
    auto B_i = hoisted_Bp.automorph(-i);
    hmm_mink->stop_computing_sc_rot(); //! SC.Rot - Stop

    // Correct Step: only 'A_i' required
    if (i) {
      hmm_mink->start_computing_sc_rot(); //! SC.Rot - Start
      auto back = hoisted_Ap.automorph(h_nrow - i);
      hmm_mink->stop_computing_sc_rot(); //! SC.Rot - Stop

      hmm_mink->start_computing_sc_scale(); //! SC.Scale - Start
      hmm_mink_cm->shift_compute_correct_mask(*A_i, *back, i);
      hmm_mink->stop_computing_sc_scale(); //! SC.Scale - Stop

      hmm_mink->start_computing_sc_add(); //! SC.Add - Start
      *A_i += *back;
      hmm_mink->stop_computing_sc_add(); //! SC.Add - Stop
    }

    hmm_mink->start_computing_sc_mult(); //! SC.Mult - Start
    A_i->multiplyBy(*B_i);
    hmm_mink->stop_computing_sc_mult(); //! SC.Mult - Stop

    hmm_mink->start_computing_sc_add(); //! SC.Add - Start
    AB += *A_i;
    hmm_mink->stop_computing_sc_add(); //! SC.Add - Stop
  }
  hmm_mink->stop_computing_sc();
  hmm_mink->sc_noise = log2_noise_bound(AB) - hmm_mink->sc_noise;
}

helib::Ctxt hhmm_mink_shift_compute_A(
    hmm_constant_multipler *cm, const helib::EncryptedArray &ea,
    const std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> &hoisted_A,
    int r) {
  auto *hmm_mink_cm = static_cast<hhmm_mink_constant_multipler *>(cm);
  ASSERT_PRINTF(hmm_mink_cm != nullptr,
                "Failed to static_cast to `hhmm_mink_constant_multipler *`\n");
  const long eh_ncol = hmm_mink_cm->eh_col();
  auto A_i = hoisted_A->automorph(-r);
  if (r) {
    auto back = hoisted_A->automorph(eh_ncol - r);
    hmm_mink_cm->shift_compute_correct_mask(*A_i, *back, r);
    *A_i += *back;
  }
  return *A_i;
}

// Single Thread Version
//* Allow A = A * B
void hhmm_mink(hmm_metric_collector *collector, hmm_constant_multipler *cm,
               const helib::EncryptedArray &ea, helib::Ctxt &AB,
               const helib::Ctxt &A, const helib::Ctxt &B) {
  auto *hmm_mink = static_cast<hmm_mink_collector *>(collector);
  ASSERT_PRINTF(hmm_mink != nullptr,
                "Failed to static_cast to `hmm_mink_collector *`\n");
  auto *hmm_mink_cm = static_cast<hhmm_mink_constant_multipler *>(cm);
  ASSERT_PRINTF(hmm_mink_cm != nullptr,
                "Failed to static_cast to `hhmm_mink_constant_multipler *`\n");

  hmm_mink->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long k = hmm_mink_cm->min_dim();
  const long h_ncol = ea.sizeOfDimension(1);
  const long eh_ncol = hmm_mink_cm->eh_col();

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);

  auto AA = A;
  auto BB = B;

  hmm_mink->start_computing_rr();
  hmm_mink->start_computing_rr_scale(); //! RA.Scale - Start
  hmm_mink_cm->clean_mask_left(AA);
  hmm_mink_cm->clean_mask_right(BB);
  hmm_mink->stop_computing_rr_scale(); //! RA.Scale - Stop
  replicate1d(hmm_mink, ea, AA, hmm_mink_cm->min_dim(), 1);
  replicate1d(hmm_mink, ea, BB, hmm_mink_cm->min_dim(), 0);
  hmm_mink->stop_computing_rr();
  hmm_mink->rr_noise = std::max(log2_noise_bound(AA), log2_noise_bound(BB)) -
                       hmm_mink->init_noise;

  //! RotateAlign - Start
  hmm_mink->ra_noise = std::max(log2_noise_bound(AA), log2_noise_bound(BB));
  hmm_mink->start_computing_ra();
  hmm_mink->start_computing_ra_rot(); //! RA.Rot - Start
  // Init rotation
  helib::GeneralAutomorphPrecon_FULL hoisted_A(AA, 1, ea);
  helib::GeneralAutomorphPrecon_FULL hoisted_B(BB, 0, ea);
  hmm_mink->stop_computing_ra_rot(); //! RA.Rot - Stop

  for (long i = 0; i < k; i++) {
    hmm_mink->start_computing_ra_rot(); //! RA.Rot - Start
    auto A_i = hoisted_A.automorph(-1l * i);
    auto B_i = hoisted_B.automorph(-1l * i);
    hmm_mink->stop_computing_ra_rot(); //! RA.Rot - Stop

    hmm_mink->start_computing_ra_scale(); //! RA.Scale - Start
    hmm_mink_cm->rotate_align_mask_left(*A_i, i);
    hmm_mink_cm->rotate_align_mask_right(*B_i, i);
    hmm_mink->stop_computing_ra_scale(); //! RA.Scale - Stop

    hmm_mink->start_computing_ra_add(); //! RA.Add - Start
    Ap += *A_i;
    Bp += *B_i;
    hmm_mink->stop_computing_ra_add(); //! RA.Add - Stop
  }
  hmm_mink->ra_noise =
      std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) - hmm_mink->rr_noise;

  //! Correct Ap (Correct RotateAlign) - Start
  hmm_mink->ra_correct_noise = log2_noise_bound(Ap);
  correct_rotate_align_A(hmm_mink, hmm_mink_cm, ea, eh_ncol, h_ncol, Ap);
  hmm_mink->ra_correct_noise =
      log2_noise_bound(Ap) - hmm_mink->ra_correct_noise;
  //! Correct Ap (Correct RotateAlign) - Completed
  hmm_mink->stop_computing_ra();
  //! RotateAlign - Completed

  //! ShiftCompute - Start
  shift_compute(hmm_mink, hmm_mink_cm, ea, eh_ncol, Ap, Bp, AB_temp);
  //! ShiftCompute - Completed

  AB = std::move(AB_temp);
  hmm_mink->final_noise = log2_noise_bound(AB);
  hmm_mink->left_cap = AB.capacity();
}

// Single Thread BSGS version
void hhmm_mink_bsgs(hmm_metric_collector *collector, hmm_constant_multipler *cm,
                    const helib::EncryptedArray &ea, helib::Ctxt &AB,
                    const helib::Ctxt &A, const helib::Ctxt &B) {
  auto *hmm_mink = static_cast<hmm_mink_collector *>(collector);
  ASSERT_PRINTF(hmm_mink != nullptr,
                "Failed to static_cast to `hmm_mink_collector *`\n");
  auto *hmm_mink_cm = static_cast<hhmm_mink_constant_multipler *>(cm);
  ASSERT_PRINTF(hmm_mink_cm != nullptr,
                "Failed to static_cast to `hhmm_mink_constant_multipler *`\n");

  hmm_mink->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long k = hmm_mink_cm->min_dim();
  const long h_ncol = ea.sizeOfDimension(1);
  const long eh_ncol = hmm_mink_cm->eh_col();

  const long f = static_cast<long>(std::ceil(std::sqrt(k)));
  const long h = (k + f - 1) / f;

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);

  auto AA = A;
  auto BB = B;

  hmm_mink->start_computing_rr();
  hmm_mink->start_computing_rr_scale(); //! RA.Scale - Start
  hmm_mink_cm->clean_mask_left(AA);
  hmm_mink_cm->clean_mask_right(BB);
  hmm_mink->stop_computing_rr_scale(); //! RA.Scale - Stop
  replicate1d(hmm_mink, ea, AA, hmm_mink_cm->min_dim(), 1);
  replicate1d(hmm_mink, ea, BB, hmm_mink_cm->min_dim(), 0);
  hmm_mink->stop_computing_rr();
  hmm_mink->rr_noise = std::max(log2_noise_bound(AA), log2_noise_bound(BB)) -
                       hmm_mink->init_noise;

  //! RotateAlign - Start
  hmm_mink->ra_noise = std::max(log2_noise_bound(AA), log2_noise_bound(BB));
  hmm_mink->start_computing_ra();
  std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);
  std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);

  hmm_mink->start_computing_ra_rot(); //! RA.Rot - Start
  // Init rotation
  helib::GeneralAutomorphPrecon_FULL hoisted_A(AA, 1, ea);
  helib::GeneralAutomorphPrecon_FULL hoisted_B(BB, 0, ea);
  hmm_mink->stop_computing_ra_rot(); //! RA.Rot - Stop

  // Pre-compute ctxts in the baby steps
  for (long i = 0; i < f; i++) {
    hmm_mink->start_computing_ra_rot(); //! RA.Rot - Start
    A_baby[i] = hoisted_A.automorph(-1l * i);
    B_baby[i] = hoisted_B.automorph(-1l * i);
    hmm_mink->stop_computing_ra_rot(); //! RA.Rot - Stop
  }

  // BSGS
  for (long i = 0; i < h; i++) {
    helib::Ctxt A_baby_sum(zero_ctxt);
    helib::Ctxt B_baby_sum(zero_ctxt);

    const long gaint_step = f * i;

    // Baby Steps
    for (long j = 0; j < f && gaint_step + j < k; j++) {
      hmm_mink->start_computing_ra_rot(); //! RA.Rot - Start
      auto A_j = *A_baby[j];
      auto B_j = *B_baby[j];
      hmm_mink->stop_computing_ra_rot(); //! RA.Rot - Stop

      hmm_mink->start_computing_ra_scale(); //! RA.Scale - Start
      hmm_mink_cm->rotate_align_mask_left(A_j, j + gaint_step);
      hmm_mink_cm->rotate_align_mask_right(B_j, j + gaint_step);
      hmm_mink->stop_computing_ra_scale(); //! RA.Scale - Stop

      hmm_mink->start_computing_ra_add(); //! RA.Add - Start
      A_baby_sum += A_j;
      B_baby_sum += B_j;
      hmm_mink->stop_computing_ra_add(); //! RA.Add - Stop
    }

    // Gaint Steps
    hmm_mink->start_computing_ra_rot(); //! RA.Rot - Start
    ea.rotate1D(A_baby_sum, 1, -1l * gaint_step);
    ea.rotate1D(B_baby_sum, 0, -1l * gaint_step);
    hmm_mink->stop_computing_ra_rot(); //! RA.Rot - Stop

    hmm_mink->start_computing_ra_add(); //! RA.Add - Start
    Ap += A_baby_sum;
    Bp += B_baby_sum;
    hmm_mink->stop_computing_ra_add(); //! RA.Add - Stop
  }
  hmm_mink->ra_noise =
      std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) - hmm_mink->rr_noise;

  //! Correct Ap (Correct RotateAlign) - Start
  hmm_mink->ra_correct_noise = log2_noise_bound(Ap);
  correct_rotate_align_A(hmm_mink, hmm_mink_cm, ea, eh_ncol, h_ncol, Ap);
  hmm_mink->ra_correct_noise =
      log2_noise_bound(Ap) - hmm_mink->ra_correct_noise;
  //! Correct Ap (Correct RotateAlign) - Completed
  hmm_mink->stop_computing_ra();
  //! RotateAlign - Completed

  //! ShiftCompute - Start
  shift_compute(hmm_mink, hmm_mink_cm, ea, eh_ncol, Ap, Bp, AB_temp);
  //! ShiftCompute - Completed

  AB = std::move(AB_temp);
  hmm_mink->final_noise = log2_noise_bound(AB);
  hmm_mink->left_cap = AB.capacity();
}

void hhmm_mink_rotate_align_A(hmm_constant_multipler *cm,
                              const helib::EncryptedArray &ea, helib::Ctxt &Ap,
                              const helib::Ctxt &A) {
  auto *hmm_mink_cm = static_cast<hhmm_mink_constant_multipler *>(cm);
  ASSERT_PRINTF(hmm_mink_cm != nullptr,
                "Failed to static_cast to `hhmm_mink_constant_multipler *`\n");

  const long k = hmm_mink_cm->min_dim();
  const long h_ncol = ea.sizeOfDimension(1);
  const long eh_ncol = hmm_mink_cm->eh_col();

  const long f = static_cast<long>(std::ceil(std::sqrt(k)));
  const long h = (k + f - 1) / f;

  helib::Ctxt zero_ctxt(A.getPubKey());
  zero_ctxt.DummyEncrypt(NTL::ZZX(0l));
  helib::Ctxt tmp(zero_ctxt);

  auto AA = A;

  hmm_mink_cm->clean_mask_left(AA);
  replicate1d(nullptr, ea, AA, hmm_mink_cm->min_dim(), 1);

  std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);

  // Init rotation
  helib::GeneralAutomorphPrecon_FULL hoisted_A(AA, 1, ea);

  // Pre-compute ctxts in the baby steps
  for (long i = 0; i < f; i++) {
    A_baby[i] = hoisted_A.automorph(-1l * i);
  }

  // BSGS
  for (long i = 0; i < h; i++) {
    helib::Ctxt A_baby_sum(zero_ctxt);
    const long gaint_step = f * i;
    // Baby Steps
    for (long j = 0; j < f && gaint_step + j < k; j++) {
      auto A_j = *A_baby[j];
      hmm_mink_cm->rotate_align_mask_left(A_j, j + gaint_step);
      A_baby_sum += A_j;
    }
    // Gaint Steps
    ea.rotate1D(A_baby_sum, 1, -1l * gaint_step);
    tmp += A_baby_sum;
  }

  //! Correct Ap (Correct RotateAlign)
  correct_rotate_align_A(nullptr, hmm_mink_cm, ea, eh_ncol, h_ncol, tmp);
  Ap = std::move(tmp);
}

void hhmm_mink_rotate_align_B(hmm_constant_multipler *cm,
                              const helib::EncryptedArray &ea, helib::Ctxt &Bp,
                              const helib::Ctxt &B) {
  auto *hmm_mink_cm = static_cast<hhmm_mink_constant_multipler *>(cm);
  ASSERT_PRINTF(hmm_mink_cm != nullptr,
                "Failed to static_cast to `hhmm_mink_constant_multipler *`\n");

  const long k = hmm_mink_cm->min_dim();
  const long h_ncol = ea.sizeOfDimension(1);
  const long eh_ncol = hmm_mink_cm->eh_col();

  const long f = static_cast<long>(std::ceil(std::sqrt(k)));
  const long h = (k + f - 1) / f;

  helib::Ctxt zero_ctxt(B.getPubKey());
  zero_ctxt.DummyEncrypt(NTL::ZZX(0l));
  helib::Ctxt tmp(zero_ctxt);

  auto BB = B;

  hmm_mink_cm->clean_mask_right(BB);
  replicate1d(nullptr, ea, BB, hmm_mink_cm->min_dim(), 0);

  std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);

  // Init rotation
  helib::GeneralAutomorphPrecon_FULL hoisted_B(BB, 0, ea);

  // Pre-compute ctxts in the baby steps
  for (long i = 0; i < f; i++) {
    B_baby[i] = hoisted_B.automorph(-1l * i);
  }

  // BSGS
  for (long i = 0; i < h; i++) {
    helib::Ctxt B_baby_sum(zero_ctxt);
    const long gaint_step = f * i;
    // Baby Steps
    for (long j = 0; j < f && gaint_step + j < k; j++) {
      auto B_j = *B_baby[j];
      hmm_mink_cm->rotate_align_mask_right(B_j, j + gaint_step);
      B_baby_sum += B_j;
    }
    // Gaint Steps
    ea.rotate1D(B_baby_sum, 0, -1l * gaint_step);
    tmp += B_baby_sum;
  }
  Bp = std::move(tmp);
}

// Multi-thread Version
void hhmm_mink_mt(hmm_metric_collector *collector, hmm_constant_multipler *cm,
                  int threads, const helib::EncryptedArray &ea, helib::Ctxt &AB,
                  const helib::Ctxt &A, const helib::Ctxt &B) {
  hmm_mink_collector *hmm_mink = static_cast<hmm_mink_collector *>(collector);
  ASSERT_PRINTF(hmm_mink != nullptr,
                "Failed to static_cast to `hmm_mink_collector *`\n");
  auto *hmm_mink_cm = static_cast<hhmm_mink_constant_multipler *>(cm);
  ASSERT_PRINTF(hmm_mink_cm != nullptr,
                "Failed to static_cast to `hhmm_mink_constant_multipler *`\n");

  hmm_mink->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long k = hmm_mink_cm->min_dim();
  const long h_ncol = ea.sizeOfDimension(1);
  const long eh_ncol = hmm_mink_cm->eh_col();

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);

  auto AA = A;
  auto BB = B;

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Ap;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Bp;
  helib::Ctxt AB_temp(zero_ctxt);
#pragma omp parallel
#pragma omp single
  {
    hmm_mink->start_computing_rr();
#pragma omp task depend(out : AA)
    hmm_mink_cm->clean_mask_left(AA);
#pragma omp task depend(out : BB)
    hmm_mink_cm->clean_mask_right(BB);
#pragma omp task depend(inout : AA)
    replicate1d(hmm_mink, ea, AA, hmm_mink_cm->min_dim(), 1);
#pragma omp task depend(inout : BB)
    replicate1d(hmm_mink, ea, BB, hmm_mink_cm->min_dim(), 0);
    hmm_mink->stop_computing_rr();
    hmm_mink->rr_noise = std::max(log2_noise_bound(AA), log2_noise_bound(BB)) -
                         hmm_mink->init_noise;

    hmm_mink->start_computing_ra(); //! RotateAlign - Start
#pragma omp task depend(in : AA)
    hoisted_A = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(AA, 1, ea);
#pragma omp task depend(in : BB)
    hoisted_B = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(BB, 0, ea);
#pragma omp taskwait //! Explicit synchronization

    std::shared_ptr<helib::Ctxt> A_i;
    std::shared_ptr<helib::Ctxt> B_i;
#pragma omp taskloop
    for (int i = 0; i < k; i++) {
      A_i = hoisted_A->automorph(-1 * i);
      B_i = hoisted_B->automorph(-1 * i);

      hmm_mink_cm->rotate_align_mask_left(*A_i, i);
      hmm_mink_cm->rotate_align_mask_right(*B_i, i);
#pragma omp critical
      {
        Ap += *A_i;
        Bp += *B_i;
      }
    }
    //! omp for: Implicit synchronization

#pragma omp task depend(out : Ap)
    {
      hmm_mink->ra_noise =
          std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) -
          hmm_mink->rr_noise;

      //! Correct Ap (Correct RotateAlign) - Start
      hmm_mink->ra_correct_noise = log2_noise_bound(Ap);
      helib::Ctxt Ap_temp = Ap;
      hmm_mink_cm->rotate_align_correct_mask(Ap_temp, Ap);
      ea.rotate1D(Ap_temp, 1l, -(h_ncol - eh_ncol));
      Ap += Ap_temp;
      hmm_mink->ra_correct_noise =
          log2_noise_bound(Ap) - hmm_mink->ra_correct_noise;
      //! Correct Ap (Correct RotateAlign) - Completed
      hmm_mink->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
      hmm_mink->stop_computing_ra(); //! RotateAlign - Completed
    }

#pragma omp task depend(in : Ap)
    {
      hmm_mink->start_computing_sc(); //! ShiftCompute - Start
      hoisted_Ap =
          std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Ap, 1, ea);
    }
#pragma omp task
    hoisted_Bp =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Bp, 0, ea);
#pragma omp taskwait //! Explicit synchronization

#pragma omp taskloop
    for (int i = 0; i < hmm_mink_cm->min_dim(); i++) {
      // Shift
      A_i = hoisted_Ap->automorph(-1l * i);
      B_i = hoisted_Bp->automorph(-1l * i);

      // Correct Step: only 'A_i'(i>0) required
      if (i) {
        auto back = hoisted_Ap->automorph(eh_ncol - i);
        hmm_mink_cm->shift_compute_correct_mask(*A_i, *back, i);
        *A_i += *back;
      }
      A_i->multiplyBy(*B_i);

#pragma omp critical
      AB_temp += *A_i;
    } //! omp taskloop: Implicit synchronization

    hmm_mink->stop_computing_sc(); //! ShiftCompute - Completed
  }
  hmm_mink->sc_noise = log2_noise_bound(AB_temp) - hmm_mink->sc_noise;

  AB = std::move(AB_temp);
  hmm_mink->final_noise = log2_noise_bound(AB);
  hmm_mink->left_cap = AB.capacity();
}

void hhmm_mink_bsgs_mt(hmm_metric_collector *collector,
                       hmm_constant_multipler *cm, int threads,
                       const helib::EncryptedArray &ea, helib::Ctxt &AB,
                       const helib::Ctxt &A, const helib::Ctxt &B) {
  hmm_mink_collector *hmm_mink = static_cast<hmm_mink_collector *>(collector);
  ASSERT_PRINTF(hmm_mink != nullptr,
                "Failed to static_cast to `hmm_mink_collector *`\n");
  auto *hmm_mink_cm = static_cast<hhmm_mink_constant_multipler *>(cm);
  ASSERT_PRINTF(hmm_mink_cm != nullptr,
                "Failed to static_cast to `hhmm_mink_constant_multipler *`\n");

  hmm_mink->init_noise = std::max(log2_noise_bound(A), log2_noise_bound(B));

  const long k = hmm_mink_cm->min_dim();
  const long h_ncol = ea.sizeOfDimension(1);
  const long eh_ncol = hmm_mink_cm->eh_col();

  const long f = static_cast<long>(std::ceil(std::sqrt(k)));
  const long h = (k + f - 1) / f;

  const helib::PubKey &pk = A.getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0));
  helib::Ctxt Ap(zero_ctxt);
  helib::Ctxt Bp(zero_ctxt);
  helib::Ctxt AB_temp(zero_ctxt);

  auto AA = A;
  auto BB = B;

  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_A;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_B;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Ap;
  std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> hoisted_Bp;

  //! RotateAlign - Start
  std::vector<std::shared_ptr<helib::Ctxt>> A_baby(f);
  std::vector<std::shared_ptr<helib::Ctxt>> B_baby(f);
#pragma omp parallel
#pragma omp single
  {
    hmm_mink->start_computing_rr();
#pragma omp task depend(out : AA)
    hmm_mink_cm->clean_mask_left(AA);
#pragma omp task depend(out : BB)
    hmm_mink_cm->clean_mask_right(BB);
#pragma omp task depend(inout : AA)
    replicate1d(hmm_mink, ea, AA, hmm_mink_cm->min_dim(), 1);
#pragma omp task depend(inout : BB)
    replicate1d(hmm_mink, ea, BB, hmm_mink_cm->min_dim(), 0);
    hmm_mink->stop_computing_rr();
    hmm_mink->rr_noise = std::max(log2_noise_bound(AA), log2_noise_bound(BB)) -
                         hmm_mink->init_noise;

    hmm_mink->start_computing_ra(); //! RotateAlign - Start
#pragma omp task depend(in : AA)
    hoisted_A = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(AA, 1, ea);
#pragma omp task depend(in : BB)
    hoisted_B = std::make_shared<helib::GeneralAutomorphPrecon_FULL>(BB, 0, ea);
#pragma omp taskwait //! Explicit synchronization

    // Pre-compute ctxts in the baby steps
#pragma omp taskloop
    for (long i = 0; i < f; i++) {
      A_baby[i] = hoisted_A->automorph(-1l * i);
      B_baby[i] = hoisted_B->automorph(-1l * i);
    } //! omp taskloop: Implicit synchronization

    // BSGS
    // FIXME: more parallelism?
#pragma omp taskloop
    for (long i = 0; i < h; i++) {
      helib::Ctxt A_baby_sum(zero_ctxt);
      helib::Ctxt B_baby_sum(zero_ctxt);

      const long gaint_step = f * i;

      // Baby Steps
      for (long j = 0; j < f && gaint_step + j < k; j++) {
        auto A_j = *A_baby[j];
        auto B_j = *B_baby[j];

        hmm_mink_cm->rotate_align_mask_left(A_j, j + gaint_step);
        hmm_mink_cm->rotate_align_mask_right(B_j, j + gaint_step);

        A_baby_sum += A_j;
        B_baby_sum += B_j;
      }

      // Gaint Steps
      ea.rotate1D(A_baby_sum, 1, -1l * gaint_step);
      ea.rotate1D(B_baby_sum, 0, -1l * gaint_step);

#pragma omp critical
      {
        Ap += A_baby_sum;
        Bp += B_baby_sum;
      }
    } //! omp taskloop: Implicit synchronization

#pragma omp task depend(out : Ap)
    {
      hmm_mink->ra_noise =
          std::max(log2_noise_bound(Ap), log2_noise_bound(Bp)) -
          hmm_mink->rr_noise;

      //! Correct Ap (Correct RotateAlign) - Start
      hmm_mink->ra_correct_noise = log2_noise_bound(Ap);
      helib::Ctxt Ap_temp = Ap;
      hmm_mink_cm->rotate_align_correct_mask(Ap_temp, Ap);
      ea.rotate1D(Ap_temp, 1, -(h_ncol - eh_ncol));
      Ap += Ap_temp;
      hmm_mink->ra_correct_noise =
          log2_noise_bound(Ap) - hmm_mink->ra_correct_noise;
      //! Correct Ap (Correct RotateAlign) - Completed
      hmm_mink->stop_computing_ra(); //! RotateAlign - Completed
      hmm_mink->sc_noise = std::max(log2_noise_bound(Ap), log2_noise_bound(Bp));
    }

    hmm_mink->start_computing_sc(); //! ShiftCompute - Start
#pragma omp task depend(in : Ap)
    hoisted_Ap =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Ap, 1, ea);
#pragma omp task
    hoisted_Bp =
        std::make_shared<helib::GeneralAutomorphPrecon_FULL>(Bp, 0, ea);
#pragma omp taskwait //! Explicit synchronization

#pragma omp taskloop
    for (int i = 0; i < hmm_mink_cm->min_dim(); i++) {
      // Shift
      auto A_i = hoisted_Ap->automorph(-i);
      auto B_i = hoisted_Bp->automorph(-i);

      // Correct Step: only 'A_i'(i>0) required
      if (i) {
        auto back = hoisted_Ap->automorph(eh_ncol - i);
        hmm_mink_cm->shift_compute_correct_mask(*A_i, *back, i);
        *A_i += *back;
      }
      A_i->multiplyBy(*B_i);

#pragma omp critical
      AB_temp += *A_i;
    } //! omp taskloop: Implicit synchronization

    hmm_mink->stop_computing_sc(); //! ShiftCompute - Completed
  }
  hmm_mink->sc_noise = log2_noise_bound(AB_temp) - hmm_mink->sc_noise;

  AB = std::move(AB_temp);
  hmm_mink->final_noise = log2_noise_bound(AB);
  hmm_mink->left_cap = AB.capacity();
}